# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：grid.py
    Created on：2025/4/8 13:51
    Author：(input)
    Description:
"""
import numpy as np
import matplotlib.pyplot as plt
import scipy.special as sp
from parasolver import ParameterSolver


class PhysicalGrid:
    def __init__(self, length, height, N, C, k):
        self.length = length
        self.height = height
        self.N = N
        self.C = C
        self.k = k
        self.k_prime = np.sqrt(1 - k**2)

        self.vertices = [
            -length/2 + 0j,
            length/2 + 0j,
            length/2 + height*1j,
            -length/2 + height*1j
        ]

        self.x = np.linspace(-length/2, length/2, N)
        self.y = np.linspace(0, height, N)
        self.X, self.Y = np.meshgrid(self.x, self.y)
        self.Z = self.X + 1j * self.Y

    def to_complex(self):
        return self.Z

    def sc_inverse(self, z):
        def sn_complex(z, k, k_prime):
            u = z.real / self.C
            v = z.imag / self.C
            sn_u, cn_u, dn_u, _ = sp.ellipj(u, k**2)
            sn_v, cn_v, dn_v, _ = sp.ellipj(v, k_prime**2)
            numerator = sn_u * dn_v + 1j * cn_u * sn_v * cn_v
            denominator = cn_v**2 + k**2 * sn_u**2 * sn_v**2
            return numerator / denominator

        w = np.zeros_like(z, dtype=complex)
        for i in range(z.shape[0]):
            for j in range(z.shape[1]):
                z_ij = z[i, j]
                w[i, j] = sn_complex(z_ij, self.k, self.k_prime)
        return w

# 辅助平面网格类（保持不变）
class AuxiliaryGrid:
    def __init__(self, W_mapped):
        self.W = W_mapped
        self.U = W_mapped.real
        self.V = W_mapped.imag

    def to_complex(self):
        return self.W

# 测试功能
def test_grids_and_transform():
    # 使用 ParameterSolver 计算 C 和 k
    length = 4.0
    height = 2.0
    solver = ParameterSolver(length, height, characteristic=1.0, dimensional=True)
    C, k, k_prime = solver.solve()

    # 创建物理平面网格
    N = 10
    phys_grid = PhysicalGrid(length, height, N, C, k)

    # 映射：物理平面 -> 辅助平面
    W_mapped = phys_grid.sc_inverse(phys_grid.Z)

    # 创建辅助平面网格（直接使用 W_mapped）
    aux_grid = AuxiliaryGrid(W_mapped)

    # 验证边界点映射
    print("验证边界点映射：")
    print(f"z = {-length/2} + 0j -> w = {phys_grid.sc_inverse(np.array([[-length/2 + 0j]]))[0, 0]}")
    print(f"z = {length/2} + 0j -> w = {phys_grid.sc_inverse(np.array([[length/2 + 0j]]))[0, 0]}")
    print(f"z = {length/2} + {height}j -> w = {phys_grid.sc_inverse(np.array([[length/2 + height*1j]]))[0, 0]}")
    print(f"z = {-length/2} + {height}j -> w = {phys_grid.sc_inverse(np.array([[-length/2 + height*1j]]))[0, 0]}")

    # 可视化
    plt.figure(figsize=(12, 4))

    # 物理平面网格
    plt.subplot(1, 2, 1)
    plt.contourf(phys_grid.X, phys_grid.Y, phys_grid.Z.real, levels=20)
    plt.plot([v.real for v in phys_grid.vertices] + [phys_grid.vertices[0].real],
             [v.imag for v in phys_grid.vertices] + [phys_grid.vertices[0].imag],
             'k-', linewidth=2, label='rectangle')
    plt.title("physical plane (z)")
    plt.xlabel("x")
    plt.ylabel("y")
    plt.axis("equal")
    plt.colorbar(label="Re(z)")
    plt.legend()

    print(aux_grid.W)
    # 辅助平面网格（映射结果：z -> w）
    plt.subplot(1, 2, 2)
    plt.contourf(aux_grid.U, aux_grid.V, aux_grid.W.real, levels=20)
    plt.title("Auxiliary plane (w, from z map)")
    plt.xlabel("u")
    plt.ylabel("v")
    plt.axis("equal")
    plt.colorbar(label="Re(w)")

    plt.tight_layout()
    plt.show()

    # 测试源汇复势（在辅助平面上）
    Q = 1.0
    source_pos = -0.5 + 0j
    sink_pos = 0.5 + 0j
    F = Q / (2 * np.pi) * (np.log(aux_grid.W - source_pos) - np.log(aux_grid.W - sink_pos))
    psi = F.imag

    # 可视化流线（辅助平面）
    plt.figure(figsize=(6, 4))
    plt.contour(aux_grid.U, aux_grid.V, psi, levels=20, colors='blue')
    plt.plot(source_pos.real, source_pos.imag, 'ro', label='source')
    plt.plot(sink_pos.real, sink_pos.imag, 'go', label='sink')
    plt.title("streamline")
    plt.xlabel("u")
    plt.ylabel("v")
    plt.axis("equal")
    plt.legend()
    plt.show()

if __name__ == "__main__":
    test_grids_and_transform()